{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ITK version:5.3.0\n",
      "C:\\Dev\\ITK-py\\Wrapping\\Generators\\Python\\itk\\__init__.py\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import itk\n",
    "print(\"ITK version:\" + itk.Version.GetITKVersion())\n",
    "print(itk.__file__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main_processing(input_image, input_labels, transform_filename, output_image, output_labels):\n",
    "    print(\"Read the input image\", input_image)\n",
    "    in_image = itk.imread(input_image)\n",
    "    print(\"Read the input labels\", input_labels)\n",
    "    labels = itk.imread(input_labels)\n",
    "    dimension = in_image.GetImageDimension()\n",
    "    transforms = itk.transformread(transform_filename)\n",
    "    direct_transform = transforms[0]\n",
    "    \n",
    "    # ApplyToImageMetadata method is not yet available in Python\n",
    "    # direct_transform.ApplyToImageMetadata(in_image)\n",
    "    # direct_transform.ApplyToImageMetadata(labels)\n",
    "    in_image = itk.resample_in_place_image_filter(in_image, rigid_transform=direct_transform)\n",
    "    labels = itk.resample_in_place_image_filter(labels, rigid_transform=direct_transform)\n",
    "    \n",
    "    print(\"Find the bounding box of the labels\")\n",
    "    bbSO = itk.ImageMaskSpatialObject[dimension].New(labels)\n",
    "    bbSO.Update()\n",
    "    \n",
    "    # this is not the tightest, because BB is computed in index space internally\n",
    "    bb = bbSO.GetMyBoundingBoxInWorldSpace()\n",
    "    start = bb.GetMinimum()\n",
    "    end = bb.GetMaximum()\n",
    "    Dimension = in_image.GetImageDimension()\n",
    "    size = [int((end[d] - start[d]) / in_image.GetSpacing()[d]) for d in range(Dimension)]\n",
    "    \n",
    "    nearest_interpolator = itk.NearestNeighborInterpolateImageFunction.New(labels)\n",
    "\n",
    "    labelsAA = itk.resample_image_filter(\n",
    "        labels,\n",
    "        interpolator=nearest_interpolator,\n",
    "        size=size,\n",
    "        output_spacing=in_image.GetSpacing(),\n",
    "        output_origin=start,\n",
    "    )\n",
    "    \n",
    "    print(\"Find the tightest bounding box of the labels\")\n",
    "    bbSO.SetImage(labelsAA);\n",
    "    bbRegion = bbSO.ComputeMyBoundingBoxInIndexSpace();\n",
    "    \n",
    "    labelsAA = itk.region_of_interest_image_filter(\n",
    "        labelsAA,\n",
    "        region_of_interest=bbRegion)\n",
    "    \n",
    "    # now resample the main image onto the same image grid\n",
    "    linear_interpolator = itk.LinearInterpolateImageFunction.New(in_image)\n",
    "    \n",
    "    out_image = itk.resample_image_filter(\n",
    "        in_image,\n",
    "        interpolator=linear_interpolator,\n",
    "        reference_image=labelsAA,\n",
    "        use_reference_image=True,\n",
    "    )\n",
    "    \n",
    "    print(\"Write the axis aligned image\", output_image)\n",
    "    itk.imwrite(out_image, output_image, compression=True)\n",
    "    print(\"Write the axis aligned labels\", output_labels)\n",
    "    itk.imwrite(labelsAA, output_labels, compression=True)\n",
    "    print(\"All done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# invoke main processing for all the images\n",
    "image_names = [\n",
    "    \"901-L\",\n",
    "    \"901-R\",\n",
    "    \"902-L\",\n",
    "    \"902-R\",\n",
    "    \"906-L\",\n",
    "    \"906-R\",\n",
    "    \"907-L\",\n",
    "    \"907-R\",\n",
    "    \"908-L\",\n",
    "    \"908-R\",\n",
    "    \"915-L\",\n",
    "    \"915-R\",\n",
    "    \"916-L\",\n",
    "    \"916-R\",\n",
    "    \"917-L\",\n",
    "    \"917-R\",\n",
    "    \"918-L\",\n",
    "    \"918-R\",\n",
    "    \"F9-3wk-01-L\",\n",
    "    \"F9-3wk-01-R\",\n",
    "    \"F9-3wk-02-L\",\n",
    "    \"F9-3wk-02-R\",\n",
    "    \"F9-3wk-03-L\",\n",
    "    \"F9-3wk-03-R\",\n",
    "    \"F9-8wk-01-L\",\n",
    "    \"F9-8wk-01-R\",\n",
    "    \"F9-8wk-02-L\",\n",
    "    \"F9-8wk-02-R\",\n",
    "]\n",
    "\n",
    "root_dir = \"../\" # root dir of the repository\n",
    "for name in image_names:\n",
    "    main_processing(root_dir+\"Data/\"+name+\".nrrd\",\n",
    "                    root_dir+\"Tibias/\"+name+\"-label.nrrd\",\n",
    "                    root_dir+\"Tibias/\"+name+\"-landmarks.tfm\",\n",
    "                    root_dir+\"Tibias/\"+name+\"-AA.nrrd\",\n",
    "                    root_dir+\"Tibias/\"+name+\"-AA-label.nrrd\"\n",
    "                   )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
